import numpy as np
from os.path import dirname, join
import random
from struct import unpack


def load_data_set(prefix):
    images = load_images(prefix)
    labels = load_labels(prefix)
    return [x for x in zip(images, labels)]


def load_images(prefix):
    path = path_for(prefix + "-images-idx3-ubyte")
    with open(path, 'rb') as f:
        magic, images_count, rows, cols = unpack(">iiii", f.read(16))
        return np.fromfile(f, dtype='uint8').reshape(images_count, rows * cols) / 255.0  # each ROW is an image


def load_labels(prefix):
    path = path_for(prefix + "-labels-idx1-ubyte")
    with open(path, 'rb') as f:
        magic, labels_count = unpack(">ii", f.read(8))
        labels = np.fromfile(f, dtype='uint8')
        assert np.size(labels) == labels_count
        return vectorize(labels, labels_count)


def vectorize(labels, labels_count):
    result = np.zeros((labels_count, 10))
    for i in range(0, labels_count):
        label = labels[i]
        result[i, label] = 1.0
    return result


def path_for(filename):
    return join(dirname(__file__), filename)


class Mnist:
    def __init__(self):
        self.training_set = load_data_set("train")
        test_set = load_data_set("t10k")
        random.shuffle(test_set)
        size = len(test_set)
        self.test_set = test_set[0:int(size/2)]
        self.dev_set = test_set[int(size/2):]

    def random_visualize(self, data_set="training_set"):
        # import matplotlib
        # matplotlib.use('qt5agg')
        import matplotlib.pyplot as plt
        data_set = getattr(self, data_set)
        images_with_label = random.sample(data_set, 25)
        fig, axes = plt.subplots(5, 5)
        i = 0
        for pair in images_with_label:
            axis = axes[int(i / 5), int(i % 5)]
            pixels = (pair[0].reshape(28, 28) * 255).astype('uint8')
            axis.title.set_text('Label is {label}'.format(label=np.argwhere(pair[1])[0, 0]))
            axis.imshow(pixels, cmap='gray')
            i += 1
        plt.show()
